import tensorflow as tf
import numpy as np
import os
from PIL import Image

os.environ["CUDA_VISIBLE_DEVICES"]="0"

im1 = Image.open('before.bmp')
im2 = Image.open('after.bmp')

x_data = np.asarray(im2)
x_data = x_data.reshape((81920,1))
a = np.loadtxt('D:/NUC.txt')
noise = np.asarray(a)
noise = noise.reshape((81920,1))
noise = noise/10
y_data = x_data.reshape((81920,1))+noise

xs = tf.placeholder(tf.float32, [None, 1])
ys = tf.placeholder(tf.float32, [None, 1])

def add_layer(inputs, in_size, out_size, activation_function=None):
  weights = tf.Variable(tf.random_normal([in_size, out_size]),name = "weights")
  biases = tf.Variable(tf.zeros([1, out_size])+0.1,name = "biases")
  Wx_plus_b = tf.matmul(inputs, weights) + biases
  

  
  if activation_function is None:
    outputs = Wx_plus_b
  else:
    outputs = activation_function(Wx_plus_b)
  return outputs

h1 = add_layer(xs, 1 , 20, activation_function=tf.nn.relu)
prediction = add_layer(h1, 20, 1, activation_function=None)

loss =  tf.reduce_mean(tf.reduce_sum(tf.square(ys-prediction),reduction_indices=[1]))
loss =  loss/(320*256)
train_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss)

init = tf.global_variables_initializer()
saver = tf.train.Saver()
sess = tf.Session()
sess.run(init)

for i in range(5000):
  sess.run(train_step, feed_dict={xs: x_data, ys: y_data})

new = sess.run(prediction, feed_dict={xs: x_data, ys: y_data})   
np.savetxt("new.txt", new); 
